{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fd72c2aa",
   "metadata": {},
   "source": [
    "# NeuroMANCER Node and System classes and modules tutorial\n",
    "\n",
    "This script demonstrates how to use NeuroMANCER Node to wrap arbitrary callable\n",
    "into symbolic representation that can be used in NeuroMANCER problem formulation.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84a23c9a",
   "metadata": {},
   "source": [
    "### Install (Colab only)\n",
    "Skip this step when running locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e418a0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install neuromancer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22e70f37",
   "metadata": {},
   "source": [
    "*Note: When running on Colab, one might encounter a pip dependency error with Lida 0.0.10. This can be ignored*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7a60398",
   "metadata": {},
   "source": [
    "### Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "26f20c73",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from neuromancer.system import Node, System"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "343f6425",
   "metadata": {},
   "source": [
    "## Node\n",
    "\n",
    "**Node** is a simple class to create symbolic modules out of arbitrary PyTorch callables.\n",
    "Node class is wrapping the callable and defines the computational node based \n",
    "on input_keys and output_keys that define computational node connections through \n",
    "intermediate dictionaries. Complex symbolic architectures can be constructed by connecting\n",
    "input and output keys of a set of Nodes via System and Problem classes.\n",
    "   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "32bde295",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['x1']\n",
      "['y1']\n",
      "{'y1': tensor([-0.6689], grad_fn=<AddBackward0>)}\n"
     ]
    }
   ],
   "source": [
    "# 1, wrap nn.Linear into Node\n",
    "net_1 = torch.nn.Linear(1, 1)\n",
    "node_1 = Node(net_1, ['x1'], ['y1'])\n",
    "# print input and output keys\n",
    "print(node_1.input_keys)\n",
    "print(node_1.output_keys)\n",
    "# evaluate forward pass of the node with dictionary input dataset\n",
    "print(node_1({'x1': torch.rand(1)}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "41bf103b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'y2': tensor([-0.3779], grad_fn=<AddBackward0>)}\n"
     ]
    }
   ],
   "source": [
    "# 2, wrap nn.Sequential into Node\n",
    "net_2 = torch.nn.Sequential(torch.nn.Linear(2, 5),\n",
    "                            torch.nn.ReLU(),\n",
    "                            torch.nn.Linear(5, 3),\n",
    "                            torch.nn.ReLU(),\n",
    "                            torch.nn.Linear(3, 1))\n",
    "node_2 = Node(net_2, ['x2'], ['y2'])\n",
    "# evaluate forward pass of the node with dictionary input dataset\n",
    "print(node_2({'x2': torch.rand(2)}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f5390740",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'y3': tensor([0.1539, 0.3889])}\n"
     ]
    }
   ],
   "source": [
    "# 3, wrap arbitrary callable into Node - allows for unwrapping the inputs\n",
    "fun_1 = lambda x1, x2: 2.*x1 - x2**2\n",
    "node_3 = Node(fun_1, ['y1', 'y2'], ['y3'], name='quadratic')\n",
    "# evaluate forward pass of the node with dictionary input dataset\n",
    "print(node_3({'y1': torch.rand(2), 'y2': torch.rand(2)}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7c9667a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'x1^2': tensor([0.0555, 0.5139]), 'x2^2': tensor([0.2789, 0.3003])}\n"
     ]
    }
   ],
   "source": [
    "# 4, wrap callable with multiple inputs and outputs\n",
    "def fun_2(x1, x2):\n",
    "    return x1**2, x2**2\n",
    "node_4 = Node(fun_2, ['x1', 'x2'], ['x1^2', 'x2^2'], name='square')\n",
    "# evaluate forward pass of the node with dictionary input dataset\n",
    "print(node_4({'x1': torch.rand(2), 'x2': torch.rand(2)}))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "972aa3d4",
   "metadata": {},
   "source": [
    "## Modules\n",
    "\n",
    "NeuroMANCER also provides implementation of useful building blocks for\n",
    "creating custom neural architectures. These include:\n",
    "* modules.blocks          - neural architecures\n",
    "* modules.activations     - custom activation functions    \n",
    "* modules.functions       - useful callables \n",
    "* modules.gnn             - graph neural nets\n",
    "* modules.rnn             - recurent neural nets\n",
    "* modules.solvers         - iterative solvers for constrained optimization\n",
    "* slim.linear             - linear algebra factorizations for weights\n",
    "        \n",
    "Next set of example shows how to wrap NeuroMANCER modules into Node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9bbaee06",
   "metadata": {},
   "outputs": [],
   "source": [
    "from neuromancer.modules import blocks\n",
    "from neuromancer.modules import activations\n",
    "from neuromancer import slim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0817c98f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'mlp': <class 'neuromancer.modules.blocks.MLP'>, 'mlp_dropout': <class 'neuromancer.modules.blocks.MLPDropout'>, 'mlp_bounds': <class 'neuromancer.modules.blocks.MLP_bounds'>, 'rnn': <class 'neuromancer.modules.blocks.RNN'>, 'pytorch_rnn': <class 'neuromancer.modules.blocks.PytorchRNN'>, 'linear': <class 'neuromancer.modules.blocks.Linear'>, 'residual_mlp': <class 'neuromancer.modules.blocks.ResMLP'>, 'basislinear': <class 'neuromancer.modules.blocks.BasisLinear'>, 'poly2': <class 'neuromancer.modules.blocks.Poly2'>, 'bilinear': <class 'neuromancer.modules.blocks.BilinearTorch'>, 'icnn': <class 'neuromancer.modules.blocks.InputConvexNN'>, 'pos_def': <class 'neuromancer.modules.blocks.PosDef'>}\n",
      "{'softexp': <class 'neuromancer.modules.activations.SoftExponential'>, 'blu': <class 'neuromancer.modules.activations.BLU'>, 'aplu': <class 'neuromancer.modules.activations.APLU'>, 'prelu': <class 'neuromancer.modules.activations.PReLU'>, 'pelu': <class 'neuromancer.modules.activations.PELU'>, 'relu': <class 'torch.nn.modules.activation.ReLU'>, 'gelu': <class 'torch.nn.modules.activation.GELU'>, 'rrelu': <class 'torch.nn.modules.activation.RReLU'>, 'hardtanh': <class 'torch.nn.modules.activation.Hardtanh'>, 'relu6': <class 'torch.nn.modules.activation.ReLU6'>, 'sigmoid': <class 'torch.nn.modules.activation.Sigmoid'>, 'hardsigmoid': <class 'torch.nn.modules.activation.Hardsigmoid'>, 'tanh': <class 'torch.nn.modules.activation.Tanh'>, 'hardswish': <class 'torch.nn.modules.activation.Hardswish'>, 'elu': <class 'torch.nn.modules.activation.ELU'>, 'celu': <class 'torch.nn.modules.activation.CELU'>, 'selu': <class 'torch.nn.modules.activation.SELU'>, 'hardshrink': <class 'torch.nn.modules.activation.Hardshrink'>, 'leakyrelu': <class 'torch.nn.modules.activation.LeakyReLU'>, 'logsigmoid': <class 'torch.nn.modules.activation.LogSigmoid'>, 'softplus': <class 'torch.nn.modules.activation.Softplus'>, 'softshrink': <class 'torch.nn.modules.activation.Softshrink'>, 'softsign': <class 'torch.nn.modules.activation.Softsign'>, 'tanhshrink': <class 'torch.nn.modules.activation.Tanhshrink'>, 'smoothedrelu': <class 'neuromancer.modules.activations.SmoothedReLU'>}\n"
     ]
    }
   ],
   "source": [
    "# for a full list of available blocks (nn.Modules) in NeuroMANCER see:\n",
    "print(blocks.blocks)\n",
    "# for a full list of available activations in NeuroMANCER see:\n",
    "print(activations.activations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "06e573e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['y3'])\n",
      "torch.Size([10, 3])\n"
     ]
    }
   ],
   "source": [
    "# 1, instantiate 4-layer multilayer perceptron with linear weight and ReLU activation\n",
    "block_1 = blocks.MLP(insize=2, outsize=3,\n",
    "                  bias=True,\n",
    "                  linear_map=slim.maps['linear'],\n",
    "                  nonlin=torch.nn.ReLU,\n",
    "                  hsizes=[80] * 4)\n",
    "# wrap modules into Node\n",
    "node_4 = Node(block_1, ['x3'], ['y3'])\n",
    "# evaluate forward pass of the node with dictionary input dataset\n",
    "data = {'x3': torch.rand(10, 2)}\n",
    "print(node_4(data).keys())\n",
    "print(node_4(data)['y3'].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "274fef37",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['y4'])\n",
      "torch.Size([10, 2])\n"
     ]
    }
   ],
   "source": [
    "# 2, instantiate recurrent neural net without bias, SVD linear map, and BLU activation\n",
    "block_2 = blocks.RNN(insize=2, outsize=2,\n",
    "                  bias=False,\n",
    "                  linear_map=slim.linear.SVDLinear,\n",
    "                  nonlin=activations.BLU,\n",
    "                  hsizes=[80] * 4)\n",
    "# wrap modules into Node\n",
    "node_5 = Node(block_2, ['x4'], ['y4'])\n",
    "# evaluate forward pass of the node with dictionary input dataset\n",
    "data = {'x4': torch.rand(10, 2)}\n",
    "print(node_5(data).keys())\n",
    "print(node_5(data)['y4'].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5d21274",
   "metadata": {},
   "source": [
    "## System\n",
    "\n",
    "**System** is a class that supports construction of cyclic computational graphs in NeuroMANCER.\n",
    "System's graph is defined by a list of Nodes. Instantiated System can be used to simulate\n",
    "dynamical systems in open or closed loop rollouts by specifying number of steps via nsteps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "13555800",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['x1', 'x2']\n",
      "['y1', 'y2', 'y3']\n",
      "{'x1': tensor([[[0.8118],\n",
      "         [0.9097],\n",
      "         [0.7591]],\n",
      "\n",
      "        [[0.1526],\n",
      "         [0.8383],\n",
      "         [0.1637]]]), 'x2': tensor([[[0.3108, 0.3442],\n",
      "         [0.8555, 0.9287],\n",
      "         [0.4466, 0.6320]],\n",
      "\n",
      "        [[0.7797, 0.3811],\n",
      "         [0.1139, 0.8550],\n",
      "         [0.7192, 0.5203]]]), 'y1': tensor([[[-0.9818],\n",
      "         [-1.0222],\n",
      "         [-0.9601]],\n",
      "\n",
      "        [[-0.7099],\n",
      "         [-0.9928],\n",
      "         [-0.7145]]], grad_fn=<CatBackward0>), 'y2': tensor([[[-0.3476],\n",
      "         [-0.4056],\n",
      "         [-0.3570]],\n",
      "\n",
      "        [[-0.3911],\n",
      "         [-0.3232],\n",
      "         [-0.3838]]], grad_fn=<CatBackward0>), 'y3': tensor([[[-2.0845],\n",
      "         [-2.2089],\n",
      "         [-2.0476]],\n",
      "\n",
      "        [[-1.5727],\n",
      "         [-2.0900],\n",
      "         [-1.5762]]], grad_fn=<CatBackward0>)}\n"
     ]
    }
   ],
   "source": [
    "# 1, create acyclic symbolic graph\n",
    "# list of nodes to construct the graph\n",
    "nodes = [node_1, node_2, node_3]\n",
    "# n steps rollout\n",
    "nsteps = 3\n",
    "# connecting nodes via System class\n",
    "system_1 = System(nodes, nsteps=nsteps)\n",
    "# print input and output keys\n",
    "print(system_1.input_keys)\n",
    "print(system_1.output_keys)\n",
    "# evaluate forward pass of the System with 3D input dataset\n",
    "batch = 2\n",
    "print(system_1({'x1': torch.rand(batch, nsteps, 1),\n",
    "                'x2': torch.rand(batch, nsteps, 2)}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "208a9928",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize symbolic computational graph\n",
    "system_1.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2dbe87c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['y1', 'x1', 'x2']\n",
      "['y1', 'y2']\n",
      "{'x1': tensor([[[0.1347],\n",
      "         [0.8139],\n",
      "         [0.4466]],\n",
      "\n",
      "        [[0.9357],\n",
      "         [0.6498],\n",
      "         [0.0123]]]), 'x2': tensor([[[0.2463, 0.3173],\n",
      "         [0.0356, 0.2024],\n",
      "         [0.6412, 0.8192]],\n",
      "\n",
      "        [[0.8419, 0.6462],\n",
      "         [0.4614, 0.4508],\n",
      "         [0.6196, 0.5821]]]), 'y1': tensor([[[-0.7025],\n",
      "         [-1.5219],\n",
      "         [-0.9827],\n",
      "         [-3.1484],\n",
      "         [-0.8312],\n",
      "         [-2.1093]],\n",
      "\n",
      "        [[-1.0330],\n",
      "         [-2.2253],\n",
      "         [-0.9150],\n",
      "         [-4.5804],\n",
      "         [-0.6520],\n",
      "         [-1.9697]]], grad_fn=<CatBackward0>), 'y2': tensor([[[-0.3419],\n",
      "         [-0.3235],\n",
      "         [-0.3794]],\n",
      "\n",
      "        [[-0.3992],\n",
      "         [-0.3605],\n",
      "         [-0.3737]]], grad_fn=<CatBackward0>)}\n"
     ]
    }
   ],
   "source": [
    "# 2, close the loop by creating recursion in one of the nodes\n",
    "nodes[2].output_keys = ['y1']\n",
    "# create new system with cyclic computational graph\n",
    "system_2 = System(nodes, nsteps=nsteps)\n",
    "# print input and output keys\n",
    "print(system_2.input_keys)\n",
    "print(system_2.output_keys)\n",
    "# evaluate forward pass of the System with 3D input dataset\n",
    "print(system_1({'x1': torch.rand(batch, nsteps, 1),\n",
    "                'x2': torch.rand(batch, nsteps, 2)}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f87c70a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize symbolic computational graph\n",
    "system_2.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0743d582",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acb914d1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "neuromancer",
   "language": "python",
   "name": "neuromancer"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
